import numpy as np
import math
from scipy.spatial.transform import Rotation
import ipdb

EPSINON = 1E-4

def rot_matrix(theta, ax):
    m = np.diag([1.0, 1.0, 1.0, 1.0])
    cos_theta = math.cos(theta)
    sin_theta = math.sin(theta)
    if(ax in ['x', 'y', 'z']):
        if(ax == 'z'):
            m[0,0] = cos_theta
            m[0,1] = -sin_theta
            m[1,0] = sin_theta
            m[1,1] = cos_theta
        if(ax == 'x'):
            m[1,1] = cos_theta
            m[1,2] = -sin_theta
            m[3,1] = sin_theta
            m[3,2] = cos_theta
        if(ax == 'y'):
            m[0,0] = cos_theta
            m[0,2] = sin_theta
            m[2,0] = -sin_theta
            m[2,2] = cos_theta
    else:
        raise Exception("wrong ax")
    return m

def trans_matrix(dis):
    assert len(dis) >= 3
    m = np.diag([1.0, 1.0, 1.0, 1.0])
    m[0,3] = dis[0]
    m[1,3] = dis[1]
    m[2,3] = dis[2]
    return m

def quatNorm(quat):
    assert len(quat) >= 4
    quat = quat[0:4]
    norm = np.sqrt(pow(quat[0], 2)+pow(quat[1], 2)+pow(quat[2], 2)+pow(quat[3], 2))
    return norm

def quaternionDotProduct(quat1, quat2):
    result = (quat1[0] * quat2[0]) + (quat1[1] * quat2[1]) + (quat1[2] * quat2[2]) + (quat1[3] * quat2[3])
    return result

def quaternionProduct(quat1, quat2):
    qx = (quat1[3] * quat2[0]) + (quat2[3] * quat1[0]) + (quat1[1] * quat2[2]) - (quat2[1] * quat1[2])
    qy = (quat1[3] * quat2[1]) + (quat2[3] * quat1[1]) + (quat1[2] * quat2[0]) - (quat2[2] * quat1[0])
    qz = (quat1[3] * quat2[2]) + (quat2[3] * quat1[2]) + (quat1[0] * quat2[1]) - (quat2[0] * quat1[1])
    qw = (quat1[3] * quat2[3]) - (quat1[0] * quat2[0]) - (quat1[1] * quat2[1]) - (quat1[2] * quat2[2])    
    quat = [qx, qy, qz, qw]
    return quat

def quaternionAdd(quat1, quat2):
    qx = quat1[0] + quat2[0]
    qy = quat1[1] + quat2[1]
    qz = quat1[2] + quat2[2]
    qw = quat1[3] + quat2[3]
    quat = [qx, qy, qz, qw]
    return quat

def quaternionMultConstant(quat1, multiplier):
    qx = quat1[0] * multiplier
    qy = quat1[1] * multiplier
    qz = quat1[2] * multiplier
    qw = quat1[3] * multiplier
    quat = [qx, qy, qz, qw]
    return quat

def quaternionSlerp(quat1, quat2, t):
    cosa = quaternionDotProduct(quat1, quat2)
    if cosa < 0.0:
        quat2 = [-1.0 * q for q in quat2]
        cosa = -1.0 * cosa
        
    scale_left = 0.0
    scale_right = 0.0
    if (abs(cosa) > (1 - EPSINON)):
        scale_left = 1 - t
        scale_right = t
    else:
        sina = math.sqrt(1.0 - cosa * cosa)
        a = np.arctan2(sina, cosa)
        scale_left = (np.sin((1.0 - t) * a)) / sina
        scale_right = np.sin(t * a) / sina
    left = quaternionMultConstant(quat1, scale_left)
    right = quaternionMultConstant(quat2, scale_right)
    quat = quaternionAdd(left, right)
    return quat

def quat2matrix(quat):
    assert len(quat) >= 4
    matrix = np.diag([1.0, 1.0, 1.0, 1.0])
    quat = quat[0:4]
    norm = math.sqrt(math.pow(quat[0], 2)+math.pow(quat[1], 2)+math.pow(quat[2], 2)+math.pow(quat[3], 2))
    qx, qy, qz, qw = [q / norm for q in quat]
    matrix[0, 0] = 1 - 2*(math.pow(qy, 2)+math.pow(qz, 2))
    matrix[0, 1] = 2*(qx*qy-qz*qw)
    matrix[0, 2] = 2*(qx*qz+qy*qw)
    matrix[1, 0] = 2*(qx*qy+qz*qw)
    matrix[1, 1] = 1 - 2*(math.pow(qx, 2)+math.pow(qz, 2))
    matrix[1, 2] = 2*(qy*qz-qx*qw)
    matrix[2, 0] = 2*(qx*qz-qy*qw)
    matrix[2, 1] = 2*(qy*qz+qx*qw)
    matrix[2, 2] = 1 - 2*(math.pow(qx, 2)+math.pow(qy, 2))
    return matrix

def matrix2quat(matrix):
    rotation_matrix = matrix[:3, :3]
    r11, r12, r13 = rotation_matrix[0]
    r21, r22, r23 = rotation_matrix[1]
    r31, r32, r33 = rotation_matrix[2]
    if (r11 + r22 + r33) >= 0.0:
        qw = 0.5 * np.sqrt(1 + r11 + r22 + r33)
        qx = 0.5 * (r32 - r23) / (2*qw)
        qy = 0.5 * (r13 - r31) / (2*qw)
        qz = 0.5 * (r21 - r12) / (2*qw)
    elif (r11 - r22 - r33) >= 0.0:
        qx = 0.5 * np.sqrt(1 + r11 - r22 - r33)
        qw = 0.5 * (r32 - r23) / (2*qx)
        qy = 0.5 * (r12 + r21) / (2*qx)
        qz = 0.5 * (r31 + r13) / (2*qx)
    elif (-r11 + r22 - r33) >= 0.0:
        qy = 0.5 * np.sqrt(1 - r11 + r22 - r33)
        qw = 0.5 * (r13 - r31) / (2*qy)
        qx = 0.5 * (r12 + r21) / (2*qy)
        qz = 0.5 * (r23 + r32) / (2*qy)
    elif (-r11 - r22 + r33) >= 0.0:
        qz = 0.5 * np.sqrt(1 - r11 - r22 + r33)
        qw = 0.5 * (r21 - r12) / (2*qz)
        qx = 0.5 * (r13 + r31) / (2*qz)
        qy = 0.5 * (r23 + r32) / (2*qz)
    else:
        raise Exception("wrong in matrix2quat function")
        
    quat = np.array([qx, qy, qz, qw])
    return quat

def matrix2Euler(matrix):
    sy = np.sqrt(matrix[0, 0] * matrix[0, 0] + matrix[1, 0] * matrix[1, 0]);
    singular = 1 if sy < EPSINON else 0
    if(singular):
        rx_angle = np.arctan2(-matrix[1, 2], matrix[1, 1])
        ry_angle = np.arctan2(-matrix[2, 0], sy)
        rz_angle = 0
    else:
        rx_angle = np.arctan2(matrix[2, 1], matrix[2, 2])
        ry_angle = np.arctan2(-matrix[2, 0], sy)
        rz_angle = np.arctan2(matrix[1, 0], matrix[0, 0])
    return [rx_angle, ry_angle, rz_angle]

class Pose:
    def __init__(self, x, y, z, qx, qy, qz, qw):
        self.x = x
        self.y = y
        self.z = z
        self.qx = qx
        self.qy = qy
        self.qz = qz
        self.qw = qw
        
    def pose_to_list(self):
        return [self.x, self.y, self.z, self.qx, self.qy, self.qz, self.qw]
    
    def pose_to_matrix(self):
        matrix = quat2matrix([self.qx, self.qy, self.qz, self.qw])
        matrix[0:3, 3] = [self.x, self.y, self.z]
        return matrix
    
    def pose_to_euler(self):
        matrix = quat2matrix([self.qx, self.qy, self.qz, self.qw])
        return matrix2Euler(matrix)
        

class DH:
    """use modified D-H to generate robotArm kinematics.
    """
    def __init__(self, alpha, a, init_theta, d):
        self.alpha = alpha  #Link twist
        self.a = a          #link length
        self.init_theta = init_theta  #joint angle
        self.d = d          #joint offset
        
    def to_matrix(self, theta):
        cos_theta = np.cos(self.init_theta + theta)
        sin_theta = np.sin(self.init_theta + theta)
        # cos_theta = np.cos(theta)
        # sin_theta = np.sin(theta)
        cos_alpha = np.cos(self.alpha)
        sin_alpha = np.sin(self.alpha)
        matrix = np.array([
                            [cos_theta,             -sin_theta,             0,              self.a              ],
                            [sin_theta*cos_alpha,   cos_theta*cos_alpha,    -sin_alpha,     -self.d*sin_alpha   ],
                            [sin_theta*sin_alpha,   cos_theta*sin_alpha,    cos_alpha,      self.d*cos_alpha    ],
                            [0,                     0,                      0,              1                   ]
                        ])
        return matrix
    
if __name__ == '__main__':
    x = Pose(0,0,0,0,0,0.7071, 0.7071)
    eu = x.pose_to_euler()
    # m = x.pose_to_matrix()
    # print(matrix2quat(m))
    # q = [0,0,0.7071,0.7071]
    # m = quat2matrix(q)
    
    # print(eu)